model_name = "MyModel"


class BaseConfig:
    """
    General configurations applied to all models
    """
    num_epochs = 20
    num_batches_show_loss = 100  # Number of batchs to show loss
    # Number of batchs to check metrics on validation dataset
    num_batches_validate = 1000
    batch_size = 64  # GPU对2的幂次的batch可以发挥更佳的性能, 太大 gpu memory 支撑不下
    learning_rate = 0.0001
    num_workers = 4  # Number of workers for data loading
    num_clicked_news_a_user = 50  # Number of sampled click history for each user
    num_words_title = 20  # title words seq len
    num_words_abstract = 50  # abstract words seq len
    word_freq_threshold = 1  # 过滤词频阈值
    entity_freq_threshold = 2  # entity阈值
    entity_confidence_threshold = 0.5
    negative_sampling_ratio = 4  # K
    dropout_probability = 0.2

    # Modify the following by the output of `src/dataprocess.py`
    num_words = 1 + 70975  # vocabulary size, embedding init,  small:1+70975,large:1+101225
    num_categories = 1 + 274  # small:1+274,large:1+295
    num_entities = 1 + 12957  # small:1+12957,large:1+21842
    num_users = 1 + 50000  # small:1+50000,large:1+711222

    word_embedding_dim = 300  # word embedding dim
    category_embedding_dim = 100
    # Modify the following only if you use another dataset
    entity_embedding_dim = 100
    # For additive attention
    query_vector_dim = 200

    inner_product_click_predict = True


class MyModelConfig(BaseConfig):
    inner_product_click_predict = True
    activate_unit = True
    use_abstract_entity = True
    use_title_entity = True

    if use_abstract_entity and use_title_entity:
        dataset_attributes = {
            "news": ['category', 'subcategory', 'title_entities', 'abstract_entities'],
            "record": ['user', 'clicked_news_length']
        }
    elif use_abstract_entity:
        dataset_attributes = {
            "news": ['category', 'subcategory', 'title', 'abstract_entities'],
            "record": ['user', 'clicked_news_length']
        }
    else:
        dataset_attributes = {
            "news": ['category', 'subcategory', 'title', 'abstract'],
            "record": ['user', 'clicked_news_length']
        }

    # 多头注意力
    num_attention_heads = 15
    entity_num_attention_heads = 10
    entity_query_vector_dim = 100

    # todo 残差块
    num_residuals = 6


